{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Slice mesh to squares"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import point_cloud_utils as pcu\n",
    "import os\n",
    "import glob\n",
    "from tqdm import tqdm\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import trimesh\n",
    "\n",
    "from trimesh import transform_points"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Config:\n",
    "    apply_pose = True\n",
    "    min_range = 2.0\n",
    "    max_range = 25.0\n",
    "\n",
    "\n",
    "class KITTIOdometryDataset:\n",
    "    def __init__(self, kitti_root_dir: str, sequence: int):\n",
    "        \"\"\"Simple KITTI DataLoader to provide a ready-to-run example.\n",
    "\n",
    "        Heavily inspired in PyLidar SLAM\n",
    "        \"\"\"\n",
    "        # Config stuff\n",
    "        self.sequence = str(int(sequence)).zfill(2)\n",
    "        self.config = Config()\n",
    "        self.kitti_sequence_dir = os.path.join(kitti_root_dir, \"sequences\", self.sequence)\n",
    "        self.velodyne_dir = os.path.join(self.kitti_sequence_dir, \"velodyne/\")\n",
    "\n",
    "        # Read stuff\n",
    "        self.calibration = self.read_calib_file(os.path.join(self.kitti_sequence_dir, \"calib.txt\"))\n",
    "        self.poses = self.load_poses(os.path.join(kitti_root_dir, f\"poses/{self.sequence}.txt\"))\n",
    "        self.scan_files = sorted(glob.glob(self.velodyne_dir + \"*.bin\"))\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        if idx < 0:\n",
    "            idx = len(self.scan_files) + idx\n",
    "        elif idx > len(self.scan_files):\n",
    "            idx -= len(self.scan_files)\n",
    "        return self.scans(idx), self.poses[idx]\n",
    "\n",
    "    def get_pose(self, idx):\n",
    "        return self.poses[idx]\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.scan_files)\n",
    "\n",
    "    def scans(self, idx):\n",
    "        return self.read_point_cloud(idx, self.scan_files[idx], self.config)\n",
    "\n",
    "    def read_point_cloud(self, idx: int, scan_file: str, config: Config):\n",
    "        points = np.fromfile(scan_file, dtype=np.float32).reshape((-1, 4))[:, :-1]\n",
    "        points = points[np.linalg.norm(points, axis=1) <= config.max_range]\n",
    "        points = points[np.linalg.norm(points, axis=1) >= config.min_range]\n",
    "        points = transform_points(points, self.poses[idx], translate=True) if config.apply_pose else points\n",
    "        return points\n",
    "\n",
    "    def load_poses(self, poses_file):\n",
    "        def _lidar_pose_gt(poses_gt):\n",
    "            _tr = self.calibration[\"Tr\"].reshape(3, 4)\n",
    "            tr = np.eye(4, dtype=np.float64)\n",
    "            tr[:3, :4] = _tr\n",
    "            left = np.einsum(\"...ij,...jk->...ik\", np.linalg.inv(tr), poses_gt)\n",
    "            right = np.einsum(\"...ij,...jk->...ik\", left, tr)\n",
    "            return right\n",
    "\n",
    "        poses = pd.read_csv(poses_file, sep=\" \", header=None).values\n",
    "        n = poses.shape[0]\n",
    "        poses = np.concatenate(\n",
    "            (poses, np.zeros((n, 3), dtype=np.float32), np.ones((n, 1), dtype=np.float32)), axis=1\n",
    "        )\n",
    "        poses = poses.reshape((n, 4, 4))  # [N, 4, 4]\n",
    "        return _lidar_pose_gt(poses)\n",
    "\n",
    "    @staticmethod\n",
    "    def read_calib_file(file_path: str) -> dict:\n",
    "        calib_dict = {}\n",
    "        with open(file_path, \"r\") as calib_file:\n",
    "            for line in calib_file.readlines():\n",
    "                tokens = line.split(\" \")\n",
    "                if tokens[0] == \"calib_time:\":\n",
    "                    continue\n",
    "                # Only read with float data\n",
    "                if len(tokens) > 0:\n",
    "                    values = [float(token) for token in tokens[1:]]\n",
    "                    values = np.array(values, dtype=np.float32)\n",
    "                    # The format in KITTI's file is <key>: <f1> <f2> <f3> ...\\n -> Remove the ':'\n",
    "                    key = tokens[0][:-1]\n",
    "                    calib_dict[key] = values\n",
    "        return calib_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = KITTIOdometryDataset(kitti_root_dir=\"/home/kowalski/data/kitti\", sequence=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Classic"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "MESH_PATH = \"/home/kowalski/data/maps_for_comparison/vdb_mesh.ply\"\n",
    "MAP_NAME = 'vdb_map'\n",
    "OUT_DIR = f\"/home/kowalski/data/sliced_maps/{MAP_NAME}\"\n",
    "SAMPLE_FREQ = 40"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### For semantic kimera"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from semantic_kitti_config import labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "MESHES_DIR = \"/home/kowalski/data/maps_for_comparison/semantic_kimera/\"\n",
    "MAP_NAME = 'kimera_true_semantic'\n",
    "OUT_DIR = f\"/home/kowalski/data/sliced_maps/{MAP_NAME}\"\n",
    "SAMPLE_FREQ = 40"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "MESHES = sorted([x.split('.')[0] for x in os.listdir(MESHES_DIR)])[5:6]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not OUT_DIR:\n",
    "    os.makedirs(OUT_DIR)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "for mesh in MESHES: \n",
    "  if not os.path.exists(os.path.join(OUT_DIR, mesh)):\n",
    "    os.makedirs(os.path.join(OUT_DIR, mesh))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for mesh in MESHES:\n",
    "    box = trimesh.creation.box(extents=(50, 50, 50))\n",
    "    map = trimesh.exchange.load.load(os.path.join(MESHES_DIR, f\"{mesh}.ply\"))\n",
    "    for idx in tqdm(range(len(dataset))[::SAMPLE_FREQ]):\n",
    "        if idx <4290:\n",
    "            continue\n",
    "        pose = dataset.get_pose(idx)\n",
    "        pose_inv = np.linalg.inv(pose)\n",
    "        box.apply_transform(pose)\n",
    "        result = map.slice_plane(box.facets_origin, -box.facets_normal)\n",
    "        box.apply_transform(pose_inv)\n",
    "        trimesh.exchange.export.export_mesh(result, os.path.join(OUT_DIR,mesh,f\"slice_{str(idx)}.ply\"), file_type=\"ply\")\n",
    "        #trimesh.exchange.export.export_mesh(box, os.path.join(OUT_DIR,mesh,f\"box_{str(idx)}.ply\"), file_type=\"ply\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sample point clouds from lidar scans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = KITTIOdometryDataset(kitti_root_dir=\"/home/kowalski/data/kitti\", sequence=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "MAP_NAME = \"kimera_gt\"\n",
    "OUT_DIR = f\"/home/kowalski/data/sample_pcls/{MAP_NAME}\"\n",
    "SAMPLE_NUM = int(5e6)\n",
    "SAMPLE_RANGE = 25\n",
    "SAMPLE_FREQ = 40\n",
    "\n",
    "if not os.path.exists(OUT_DIR):\n",
    "  os.makedirs(OUT_DIR)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_points_in_range(points, pose, range):\n",
    "    inv_pose = np.linalg.inv(pose)\n",
    "    points_t = transform_points(points, inv_pose, translate=True)\n",
    "    \n",
    "    # cut points inside specidic range\n",
    "    x_min, x_max = -range, range\n",
    "    y_min, y_max = -range, range\n",
    "    z_min, z_max = -range, range\n",
    "    ll = np.array([x_min, y_min, z_min])\n",
    "    ur = np.array([x_max, y_max, z_max])\n",
    "\n",
    "    inidx = np.all(np.logical_and(ll <= points_t, points_t <= ur), axis=1)\n",
    "\n",
    "    return points[inidx]\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for idx (scan, pose) in tqdm(enumerate(dataset)):\n",
    "    if idx // SAMPLE_FREQ:\n",
    "        continue\n",
    "    merged_pcl = scan\n",
    "    for _i in range(len(dataset)):\n",
    "        if _i == idx:\n",
    "            continue\n",
    "        p = dataset.get_pose(_i)\n",
    "        if np.linalg.norm(p[:3,3] - pose[:3, 3]) < SAMPLE_RANGE:\n",
    "            s, p = dataset[_i]\n",
    "            merged_pcl = np.concatenate((merged_pcl, s), axis=0)\n",
    "        #print(len(merged_pcl))\n",
    "    points_inside = get_points_in_range(merged_pcl, pose, SAMPLE_RANGE)\n",
    "    replace = len(points_inside) < SAMPLE_NUM\n",
    "    if replace:\n",
    "        print(f\"Replace!: {len(points_inside)}\" )\n",
    "    sampled_points = points_inside[np.random.choice(range(len(points_inside)), SAMPLE_NUM, replace=replace)]\n",
    "    sampled_points = trimesh.PointCloud(sampled_points)\n",
    "    sampled_points.export(os.path.join(OUT_DIR, f\"sampled_pcl_{idx}.ply\"), file_type=\"ply\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sample point clouds meshes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "MAP_NAME = \"vdb_map\"\n",
    "MESH_PATH = f\"/home/kowalski/data/sliced_maps/{MAP_NAME}/\"\n",
    "SAMPLED_PATH = f\"/home/kowalski/data/sample_pcls/{MAP_NAME}/\"\n",
    "GT = f\"/home/kowalski/data/sample_pcls/gt/\"\n",
    "SAMPLE_NUM = int(5e6)\n",
    "OUT_DIR = f\"/home/kowalski/data/sample_pcls/{MAP_NAME}\"\n",
    "\n",
    "meshes = sorted(os.listdir(MESH_PATH), key=lambda s: int(s.split('_')[-1].split('.')[0]))\n",
    "pcl_clouds = sorted(os.listdir(GT), key=lambda s: int(s.split('_')[-1].split('.')[0]))\n",
    "sampled_clouds = sorted(os.listdir(GT), key=lambda s: int(s.split('_')[-1].split('.')[0]))\n",
    "\n",
    "if not os.path.exists(OUT_DIR):\n",
    "  os.makedirs(OUT_DIR)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "import point_cloud_utils as pcu\n",
    "import open3d as o3d"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for sample_p, pcl_p in tqdm(zip(sampled_clouds, pcl_clouds)):\n",
    "\n",
    "    idx = int(pcl_p.split('_')[-1].split('.')[0])\n",
    "    pcd = o3d.io.read_point_cloud(os.path.join(GT, pcl_p))\n",
    "    p = np.asarray(pcd.points)\n",
    "\n",
    "    pcd = o3d.io.read_point_cloud(os.path.join(SAMPLED_PATH, sample_p))\n",
    "    v = np.asarray(pcd.points)\n",
    "\n",
    "    idx = np.random.choice(range(len(p)), int(0.2* v.shape[0]))\n",
    "\n",
    "    # Use the indices to get the sample positions and normals\n",
    "    p_sampled = p[idx]\n",
    "    v_sampled = v[idx]\n",
    "    \n",
    "    base = trimesh.PointCloud(v_sampled, colors=colors_b)\n",
    "    gt = trimesh.PointCloud(p_sampled, colors=colors_g)\n",
    "    \n",
    "    d, fi, bc = pcu.closest_points_on_mesh(p, v, f)\n",
    "    closest_points = pcu.interpolate_barycentric_coords(f, fi, bc, v)\n",
    "    pcl = trimesh.PointCloud(closest_points)\n",
    "    idx = int(mesh_p.split('_')[-1].split('.')[0])\n",
    "    pcl.export(os.path.join(OUT_DIR, f\"sampled_pcl_{idx}.ply\"), file_type=\"ply\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Metrics!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import csv\n",
    "import open3d as o3d\n",
    "import numpy as np\n",
    "from metrics import CloudMetrics\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "MAP_NAME = \"poisson_map\"\n",
    "SAMPLES_PATH = f\"/home/kowalski/data/sample_pcls/{MAP_NAME}\"\n",
    "GT = f\"/home/kowalski/data/sample_pcls/gt/\"\n",
    "OUT_DIR = f\"./metrics/{MAP_NAME}\"\n",
    "\n",
    "samples_mesh = sorted(os.listdir(SAMPLES_PATH), key=lambda s: int(s.split('_')[-1].split('.')[0]))[:100]\n",
    "samples_gt = sorted(os.listdir(GT), key=lambda s: int(s.split('_')[-1].split('.')[0]))[:100]\n",
    "\n",
    "if not os.path.exists(OUT_DIR):\n",
    "  os.makedirs(OUT_DIR)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "metrics_log = {\"Hdf_fw\":[], \"Hdf_bw\":[], \"CD\":[], \"precision\":[], \"recall\": [], \"f-score\": []}\n",
    "for mesh_p, gt_p in tqdm(zip(samples_mesh, samples_gt)):\n",
    "    mesh_sample = o3d.io.read_point_cloud(os.path.join(SAMPLES_PATH, mesh_p))\n",
    "    gt_sample = o3d.io.read_point_cloud(os.path.join(GT, gt_p))\n",
    "    metrics = CloudMetrics(mesh_sample, gt_sample)\n",
    "    hausdorff_forward, hausdorff_backward = metrics.hausdorff_distance()\n",
    "    chamfer_distance = metrics.chamfer_distance()\n",
    "    P, R, F = metrics.pr_f_score(0.03)\n",
    "    metrics_log[\"Hdf_fw\"].append(hausdorff_forward)\n",
    "    metrics_log[\"Hdf_bw\"].append(hausdorff_backward)\n",
    "    metrics_log[\"CD\"].append(chamfer_distance)\n",
    "    metrics_log[\"precision\"].append(P)\n",
    "    metrics_log[\"recall\"].append(R)\n",
    "    metrics_log[\"f-score\"].append(F)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"{MAP_NAME} Evaluation:\")\n",
    "print(\"Forward Hausdorff distance: \", np.median(np.array(metrics_log[\"Hdf_fw\"])).round(5))\n",
    "print(\"Bakward Hausdorff distance: \",  np.median(np.array(metrics_log[\"Hdf_bw\"])).round(5))\n",
    "print(\"          Chamfer distance: \",  np.median(np.array(metrics_log[\"CD\"])).round(5))\n",
    "print(\"                 Precision: \", np.round( np.median(np.array(metrics_log[\"precision\"])), 5))\n",
    "print(\"                    Recall: \", np.round( np.median(np.array(metrics_log[\"recall\"])), 5))\n",
    "print(\"                   F-score: \", np.round( np.median(np.array(metrics_log[\"f-score\"])), 5))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(os.path.join(OUT_DIR, \"results.csv\"), 'w',  newline='') as csvfile:\n",
    "    csvwriter = csv.writer(csvfile)\n",
    "    csvwriter.writerow(['id',\n",
    "                        'Forward Hausdorff distance',\n",
    "                        \"Bakward Hausdorff distance\",\n",
    "                        \"Chamfer distance\",\n",
    "                        \"Precision\",\n",
    "                        \"Recall\",\n",
    "                        \"F-score\"\n",
    "                        ])\n",
    "    for idx, sample in  enumerate(samples_mesh):\n",
    "        csvwriter.writerow([str(int(sample.split('_')[-1].split('.')[0])),\n",
    "                            str(metrics_log[\"Hdf_fw\"][idx]),\n",
    "                            str(metrics_log[\"Hdf_bw\"][idx]),\n",
    "                            str(metrics_log[\"CD\"][idx]),\n",
    "                            str(metrics_log[\"precision\"][idx]),\n",
    "                            str(metrics_log[\"recall\"][idx]),\n",
    "                            str(metrics_log[\"f-score\"][idx]),\n",
    "                           ])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(os.path.join(OUT_DIR, 'metric.txt'), 'w+') as f:\n",
    "    f.write(f\"{MAP_NAME} Evaluation:\\n\")\n",
    "    f.write(f\"Forward Hausdorff distance: {np.array(metrics_log['Hdf_fw']).mean().round(5)}\\n\")\n",
    "    f.write(f\"Bakward Hausdorff distance: {np.array(metrics_log['Hdf_bw']).mean().round(5)}\\n\")\n",
    "    f.write(f\"          Chamfer distance: {np.array(metrics_log['CD']).mean().round(5)}\\n\")\n",
    "    f.write(f\"                 Precision: {np.round( np.array(metrics_log['precision']).mean(), 5)}\\n\")\n",
    "    f.write(f\"                    Recall: {np.round( np.array(metrics_log['recall']).mean(), 5)}\\n\")\n",
    "    f.write(f\"                   F-score: {np.round( np.array(metrics_log['f-score']).mean(), 5)}\\n\")"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "cea7fedf4a86d138f6cff805221b2617639e11c04df65c7f5fe3cc4127107364"
  },
  "kernelspec": {
   "display_name": "Python 3.9.12 ('vdbfusion')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.12"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
